library(shiny)
library(dplyr)
library(ggplot2)
library(grid)
## Copyright Ling Lab at UMass Chan Medical School by Pei-Lun Kao
# =========================================================
# Helper functions (ROBUST & Q2=0 SAFE Version 1.1)
# =========================================================

# Read ddPCR CSV safely
read_ddpcr <- function(path) {
  df <- read.csv(path, skip = 4, header = TRUE, stringsAsFactors = FALSE)
  
  required_cols <- c("Ch1Amplitude", "Ch2Amplitude")
  validate(
    need(all(required_cols %in% colnames(df)),
         "Amplitude columns not detected. Check CSV format.")
  )
  
  df %>%
    mutate(
      Ch1Amplitude = as.numeric(Ch1Amplitude),
      Ch2Amplitude = as.numeric(Ch2Amplitude)
    ) %>%
    filter(!is.na(Ch1Amplitude), !is.na(Ch2Amplitude))
}

# Assign quadrants
assign_quadrant <- function(df, x_th, y_th) {
  df %>%
    mutate(
      quadrant = case_when(
        Ch2Amplitude >= x_th & Ch1Amplitude >= y_th ~ "Q2",
        Ch2Amplitude <  x_th & Ch1Amplitude >= y_th ~ "Q1",
        Ch2Amplitude <  x_th & Ch1Amplitude <  y_th ~ "Q4",
        Ch2Amplitude >= x_th & Ch1Amplitude <  y_th ~ "Q3"
      )
    )
}

# Quadrant statistics
quad_stats <- function(df) {
  df %>%
    group_by(quadrant) %>%
    summarise(
      count = n(),
      mean_Ch2 = mean(Ch2Amplitude),
      mean_Ch1 = mean(Ch1Amplitude),
      .groups = "drop"
    )
}

# ---- CRITICAL FIX: Safe quadrant counts ----
safe_qcount <- function(qs, q) {
  if (q %in% qs$quadrant) {
    qs$count[qs$quadrant == q]
  } else {
    0
  }
}

# Copy concentration calculation (Q2 = 0 safe)
calc_conc <- function(qs, n_samples = 1, include_dp = TRUE) {
  
  Q1 <- safe_qcount(qs, "Q1")  # FAM
  Q2 <- safe_qcount(qs, "Q2")  # DP
  Q3 <- safe_qcount(qs, "Q3")  # HEX
  Q4 <- safe_qcount(qs, "Q4")  # NT
  
  N_tot <- Q1 + Q2 + Q3 + Q4
  
  validate(
    need(N_tot > 0, "No droplets detected")
  )
  
  V_drop <- (20 * n_samples) / N_tot
  
  if (include_dp) {
    PFAM <- Q1 + Q2
    PHEX <- Q3 + Q2
  } else {
    PFAM <- Q1
    PHEX <- Q3
  }
  
  validate(
    need(PFAM < N_tot, "FAM channel saturated"),
    need(PHEX < N_tot, "HEX channel saturated")
  )
  
  list(
    C_FAM = -(1.1 / V_drop) * log(1 - PFAM / N_tot),
    C_HEX = -(1.1 / V_drop) * log(1 - PHEX / N_tot)
  )
}

# Fractional abundance (assay‑aware & safe)
calc_fa <- function(C_FAM, C_HEX, assay) {
  
  denom <- C_FAM + C_HEX
  validate(
    need(denom > 0, "Unable to compute FA (zero total concentration)")
  )
  
  if (assay == "WT_Mut") {
    list(
      fam_val = C_FAM / denom,
      hex_val = C_HEX / denom,
      fam_label = "Mut_FA",
      hex_label = "WT_FA"
    )
  } else {
    list(
      fam_val = C_FAM / denom,
      hex_val = C_HEX / denom,
      fam_label = "Edit‑2_FA",
      hex_label = "Mut_FA"
    )
  }
}

# Plot function (legible + publication‑safe)
make_plot <- function(df, title, conc, fa, x_th, y_th) {
  ggplot(df, aes(Ch2Amplitude, Ch1Amplitude, color = quadrant)) +
    
    geom_point(size = 0.5, alpha = 0.35) +
    
    geom_hline(yintercept = y_th,
               linetype = "dashed",
               color = "blue",
               linewidth = 0.7) +
    
    geom_vline(xintercept = x_th,
               linetype = "dashed",
               color = "red",
               linewidth = 0.7) +
    
    scale_color_manual(
      values = c(
        Q1 = "#1b9e77",
        Q2 = "#7570b3",
        Q3 = "#d95f02",
        Q4 = "grey60"
      )
    ) +
    
    annotate(
      "label", x = 1500, y = 4500,
      label = paste("HEX (cp/uL) =", round(conc$C_HEX, 3)),
      size = 5, fill = "white"
    ) +
    
    annotate(
      "label", x = 1500, y = 4000,
      label = paste("FAM (cp/uL) =", round(conc$C_FAM, 3)),
      size = 5, fill = "white"
    ) +
    
    annotate(
      "label", x = 1500, y = 3500,
      label = paste(fa$fam_label, "=", round(fa$fam_val, 3)),
      size = 5, fill = "white"
    ) +
    
    annotate(
      "label", x = 1500, y = 3000,
      label = paste(fa$hex_label, "=", round(fa$hex_val, 3)),
      size = 5, fill = "white"
    ) +
    
    labs(
      title = title,
      x = "HEX Amplitude",
      y = "FAM Amplitude"
    ) +
    
    theme_classic(base_size = 14) +
    theme(
      legend.position = "right",
      legend.key.size = unit(1.5, "cm"),
      legend.title = element_blank(),
      plot.title = element_text(face = "bold", hjust = 0.5),
      axis.title = element_text(face = "bold"),
      guides(color = guide_legend(override.aes = list(shape = 21, size  = 5, alpha = 1))
             ) 
      )
}

# =========================================================
# UI
# =========================================================

ui <- fluidPage(
  titlePanel("Ling Lab ddPCR Analysis Application -- By P.L. Kao"),
  
  sidebarLayout(
    sidebarPanel(
      fileInput("files", "Upload ddPCR 'Amplitude' file export .csv formats only", multiple = TRUE, accept = ".csv"),
      
      hr(),
      selectInput("sample1_file", "Sample 1 file", choices = NULL),
      textInput("sample1_name", "Sample 1 label", "Sample1"),
      
      selectInput("sample2_file", "Sample 2 file", choices = NULL),
      textInput("sample2_name", "Sample 2 label", "Sample2"),
      
      hr(),
      numericInput("x_th", "HEX threshold (Ch2)", 3000),
      numericInput("y_th", "FAM threshold (Ch1)", 2000),
      
      radioButtons(
        "dp_mode",
        "DP droplets in target counts",
        choices = c("Include DP" = TRUE, "Exclude DP" = FALSE),
        selected = TRUE
      ),
      
      radioButtons(
        "assay_type",
        "Assay type",
        choices = c(
          "WT + Mut (FAM = Mut, HEX = WT)" = "WT_Mut",
          "Edit‑2 + Mut (FAM = Edit‑2, HEX = Mut)" = "Edit2_Mut"
        ),
        selected = "WT_Mut"
      )
    ),
    
    mainPanel(
      tabsetPanel(
        tabPanel("Sample 1", plotOutput("plot1")),
        tabPanel("Sample 2", plotOutput("plot2")),
        tabPanel("Merged",   plotOutput("plot3")),
        tabPanel("QC Stats", tableOutput("qc_table"))
      )
    )
  )
)

# =========================================================
# Server
# =========================================================

server <- function(input, output, session) {
  
  observeEvent(input$files, {
    updateSelectInput(session, "sample1_file", choices = input$files$name)
    updateSelectInput(session, "sample2_file", choices = input$files$name)
  })
  
  data_list <- reactive({
    req(input$files)
    dfs <- lapply(input$files$datapath, read_ddpcr)
    names(dfs) <- input$files$name
    dfs
  })
  
  Sample1 <- reactive({
    assign_quadrant(
      data_list()[[input$sample1_file]],
      input$x_th, input$y_th
    )
  })
  
  Sample2 <- reactive({
    assign_quadrant(
      data_list()[[input$sample2_file]],
      input$x_th, input$y_th
    )
  })
  
  Merged <- reactive({
    rbind(Sample1(), Sample2()) %>%
      assign_quadrant(input$x_th, input$y_th)
  })
  
  output$plot1 <- renderPlot({
    qs <- quad_stats(Sample1())
    conc <- calc_conc(qs, include_dp = as.logical(input$dp_mode))
    fa <- calc_fa(conc$C_FAM, conc$C_HEX, input$assay_type)
    make_plot(Sample1(), input$sample1_name, conc, fa,
              input$x_th, input$y_th)
  })
  
  output$plot2 <- renderPlot({
    qs <- quad_stats(Sample2())
    conc <- calc_conc(qs, include_dp = as.logical(input$dp_mode))
    fa <- calc_fa(conc$C_FAM, conc$C_HEX, input$assay_type)
    make_plot(Sample2(), input$sample2_name, conc, fa,
              input$x_th, input$y_th)
  })
  
  output$plot3 <- renderPlot({
    qs <- quad_stats(Merged())
    conc <- calc_conc(qs, n_samples = 2,
                      include_dp = as.logical(input$dp_mode))
    fa <- calc_fa(conc$C_FAM, conc$C_HEX, input$assay_type)
    make_plot(Merged(), "Merged Sample1 + Sample2", conc, fa,
              input$x_th, input$y_th)
  })
  
  
  output$qc_table <- renderTable({
    
    qs1 <- quad_stats(Sample1()) %>%
      mutate(Sample = input$sample1_name)
    
    qs2 <- quad_stats(Sample2()) %>%
      mutate(Sample = input$sample2_name)
    
    qs3 <- quad_stats(Merged()) %>%
      mutate(Sample = "Merged")
    
    bind_rows(qs1, qs2, qs3) %>%
      select(Sample, everything())
    
  })
  
}

shinyApp(ui, server)
``